package com.stool.gram.domain;

import com.stool.core.Classifier;
import com.stool.core.ModelClassifier;
import com.stool.pojo.ModelConfiguration;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.File;
import java.net.URL;
import java.util.HashMap;
import java.util.Map;

public class DLmodels {

    private final Map<String, Classifier> modelMap = new HashMap<>();

    private String dir;

    public DLmodels(String dir) {
        this.dir = dir;
        build();
    }

    public Classifier get(String computeType) {
        return modelMap.get(computeType);
    }

    private void build() {
        createClassifier("0-5");
        createClassifier("0-10");
        createClassifier("0-15");
        createClassifier("0-20");
        createClassifier("0-25");
        createClassifier("0-30");
        createClassifier("0-35");
        createClassifier("0-40");
    }

    private String getPath(String name) {
        URL url = DLmodels.class.getClassLoader().getResource(name);
        String modelDir = url.getPath();
        return modelDir;
    }

    private void createClassifier(String computeType) {
        ModelConfiguration configuration = new ModelConfiguration();

        String path = dir + File.separator + "lstm_" + computeType;

        configuration.setModelDir(path);
        configuration.setModelName("model_inference.pb");
        configuration.setLabelFileName("label.txt");
        configuration.setInputTensorName("x-input");
        configuration.setOutputTensorName("output");
        configuration.setNumberLabel(2);
        final Classifier classifier = ModelClassifier.getClassifier(configuration);
        modelMap.put(computeType, classifier);
    }

}
